%% Figure 3a and Extended data Figure 2
% Heat maps of theta in time (experiment)
bins = 10;
t = 6;
bin_edges = linspace(-pi, pi, bins + 1);
spacing = 5;
%---9 cp
load('./9cp/thetaAndtracks.mat')

[counts, centers] = hist(theta_norm',(bin_edges(2:end) + bin_edges(1:end-1))/2);
counts = (counts./sum(counts))./(2*pi/bins);
len = (t*30)/spacing;  %this is where you set the length in terms of frames

figure; imagesc([1 len*spacing]./30, centers, counts(:,1:len));
cmap = colormap('hot');
colormap(cmap); cbar = colorbar;
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
title('experiment');xlabel('t (s)'); ylabel('theta (radians)')
caxis([0 0.25])
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth = 1.5;
cbar.LineWidth = 1;
cbar.Color = 'k';

print(['./9cp/theta_map_exp_t' num2str(t) '.pdf'],'-dpdf','-r1200')

% main figure
% writematrix("time (s)", 'Figure3.xlsx', 'Sheet', '3A-Top', 'Range', 'B1')
% writematrix(linspace(0, len*spacing./30, len), 'Figure3.xlsx', 'Sheet', '3A-Top', 'Range', 'C1')
% writematrix("orientation (rad)", 'Figure3.xlsx', 'Sheet', '3A-Top', 'Range', 'A2')
% writematrix(centers, 'Figure3.xlsx', 'Sheet', '3A-Top', 'Range', 'A3:A13')
% writematrix(counts(:,1:len), 'Figure3.xlsx', 'Sheet', '3A-Top', 'Range', 'C3')
% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2D', 'Range', 'B1')
% writematrix(linspace(0, len*spacing./30, len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2D', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2D', 'Range', 'A2')
% writematrix(centers, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2D', 'Range', 'A3:A13')
% writematrix(counts(:,1:len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2D', 'Range', 'C3')


%---control

load('./control/thetaAndtracks.mat')

[counts, centers] = hist(theta_norm',(bin_edges(2:end) + bin_edges(1:end-1))/2);
counts = (counts./sum(counts))./(2*pi/bins);
len = (t*30)/spacing;  %this is where you set the length in terms of frames

figure; imagesc([1 len*spacing]./30, centers, counts(:,1:len)); colormap hot; colorbar
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
title('experiment');xlabel('t (s)'); ylabel('theta (radians)')
caxis([0 0.25])
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth = 1.5;

print(['./control/theta_map_exp_t' num2str(t) '.pdf'],'-dpdf','-r1200')

% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2A', 'Range', 'B1')
% writematrix(linspace(0, len*spacing./30, len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2A', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2A', 'Range', 'A2')
% writematrix(centers, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2A', 'Range', 'A3:A13')
% writematrix(counts(:,1:len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2A', 'Range', 'C3')
%----------------Heat maps of theta in time(simulation)-------------------
%
bins = 20;
t = 6;
%---9cp

load('./9cp/SimulationTracks.mat')
len = t/(.02*10); %each time step in simulation corresponds to 0.02 simulation time units (50 = 1 t)
[counts, centers] = hist(wrapToPi(Theta(1:len,:)+pi)',bins);
counts = counts./sum(counts);

temp = counts./counts(:,1); temp = (temp./sum(temp))./(2*pi/bins);
figure; imagesc(10*[0 len*.02], centers, temp);
caxis([0 0.25])

cmap = ColorMapDerek([0 0 0], [1 0 0], 'no');
cmap = colormap('hot');
colormap(cmap); cbar = colorbar;

title('simulation')
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
xlabel('$t = \tilde{t}W/u$ (s)','interpreter','latex'); ylabel('theta (radians)')
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth =1.5;

print(['./9cp/theta_map_sim_t' num2str(t) '.pdf'],'-dpdf','-r1200')

% main figure
% writematrix("time (s)", 'Figure3.xlsx', 'Sheet', '3A-Bottom', 'Range', 'B1')
% writematrix(linspace(0, 6, len), 'Figure3.xlsx', 'Sheet', '3A-Bottom', 'Range', 'C1')
% writematrix("orientation (rad)", 'Figure3.xlsx', 'Sheet', '3A-Bottom', 'Range', 'A2')
% writematrix(centers, 'Figure3.xlsx', 'Sheet', '3A-Bottom', 'Range', 'A3')
% writematrix(temp, 'Figure3.xlsx', 'Sheet', '3A-Bottom', 'Range', 'C3')
% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2E', 'Range', 'B1')
% writematrix(linspace(0, 6, len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2E', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2E', 'Range', 'A2')
% writematrix(centers, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2E', 'Range', 'A3')
% writematrix(temp, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2E', 'Range', 'C3')
%
%---control

load('./control/SimulationTracks.mat')
len = t/(.02*10); %each time step in simulation corresponds to 0.02 simulation time units (50 = 1 t)
[counts, centers] = hist(wrapToPi(Theta(1:len,:)+pi)',bins);
counts = counts./sum(counts);

temp = counts./counts(:,1); temp = (temp./sum(temp))./(2*pi/bins);
figure; imagesc(10*[0 len*.02], centers, temp);
caxis([0 0.25])
colormap hot; colorbar; title('simulation')
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
xlabel('$t = \tilde{t}W/u$ (s)','interpreter','latex'); ylabel('theta (radians)')
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth = 1.5;

print(['./control/theta_map_sim_t' num2str(t) '.pdf'],'-dpdf','-r1200')

% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2B', 'Range', 'B1')
% writematrix(linspace(0, 6, len), 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2B', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2B', 'Range', 'A2')
% writematrix(centers, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2B', 'Range', 'A3')
% writematrix(temp, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2B', 'Range', 'C3')

%---9cp---theory
theta_x_theor = linspace(-pi, pi, 501);
t = 6;
tau = 14.52;
Drot = 0.069;
p = FPdistribution(theta_x_theor, 0:0.001:t, Drot, tau);

figure; imagesc(0:0.001:t, theta_x_theor, p');
caxis([0 0.25])
colormap hot; colorbar; title('theory')
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
xlabel('$t = \tilde{t}W/u$ (s)','interpreter','latex'); ylabel('theta (radians)')
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth = 1.5;

print(['./9cp/theta_map_theory.pdf'],'-dpdf','-r1200')

% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2F', 'Range', 'B1')
% writematrix([0:0.001:t], 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2F', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2F', 'Range', 'A2')
% writematrix(theta_x_theor', 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2F', 'Range', 'A3')
% writematrix(p', 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2F', 'Range', 'C3')

%---control---theory
theta_x_theor = linspace(-pi, pi, 501);
t = 6;
tau = 1000;
Drot = 0.069;
p = FPdistribution(theta_x_theor, 0:0.001:t, Drot, tau);

figure; imagesc(0:0.001:t, theta_x_theor, p');
caxis([0 0.25])
colormap hot; colorbar; title('theory')
yticks([-pi -pi/2 0 pi/2 pi])
yticklabels({});
xlabel('$t = \tilde{t}W/u$ (s)','interpreter','latex'); ylabel('theta (radians)')
ax = gca;
ax.XColor = 'k';
ax.YColor = 'k';
ax.LineWidth = 1.5;

print(['./control/theta_map_theory.pdf'],'-dpdf','-r1200')

% sup figure
% writematrix("time (s)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2C', 'Range', 'B1')
% writematrix(0:0.001:t, 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2C', 'Range', 'C1')
% writematrix("orientation (rad)", 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2C', 'Range', 'A2')
% writematrix(theta_x_theor', 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2C', 'Range', 'A3')
% writematrix(p', 'ExtDataFigure2.xlsx', 'Sheet', 'Ext2C', 'Range', 'C3')
%% FUNCTION

function p = FPdistribution(theta, t, Dr, tau)
Nt = length(t);
Ntheta = length(theta);
dt = abs(t(2)-t(1));
dtheta = abs(theta(2)-theta(1));
p0 = 1/2/pi*ones(size(theta));

p = zeros(Nt,Ntheta);

for ii = 1:Nt
    
    if ii == 1
        p(ii,:) = p0;
    else
        p(ii,2:end-1) =  p(ii-1,2:end-1) ...
            + dt*(   Dr * (p(ii-1,3:end) - 2*p(ii-1,2:end-1) + p(ii-1,1:end-2))/dtheta^2 ...
            + (1/tau) * (p(ii-1,3:end).*sin(theta(3:end)) - p(ii-1,1:end-2).*sin(theta(1:end-2)))/2/dtheta);
        
        p(ii,1) =  p(ii-1,1) ...
            + dt*(   Dr * (p(ii-1,2) - 2*p(ii-1,1) + p(ii-1,end))/dtheta^2 ...
            + (1/tau) * (p(ii-1,2).*sin(theta(2)) - p(ii-1,end).*sin(theta(end)))/2/dtheta);
        p(ii,end) =  p(ii-1,end) ...
            + dt*(   Dr * (p(ii-1,1) - 2*p(ii-1,end) + p(ii-1,end-1))/dtheta^2 ...
            + (1/tau) * (p(ii-1,1).*sin(theta(1)) - p(ii-1,end-1).*sin(theta(end-1)))/2/dtheta);
    end
    
    
end
end
%% FUNCTION
function [param,stat]=sigm_fit(x,y,fixed_params,initial_params,plot_flag)
% Optimization of parameters of the sigmoid function
%
% Syntax:
%       [param]=sigm_fit(x,y)
%
%       that is the same that
%       [param]=sigm_fit(x,y,[],[],[])     % no fixed_params, automatic initial_params
%
%       [param]=sigm_fit(x,y,fixed_params)        % automatic initial_params
%       [param]=sigm_fit(x,y,[],initial_params)   % use it when the estimation is poor
%       [param]=sigm_fit(x,y,fixed_params,initial_params,plot_flag)
%
% param = [min, max, x50, slope]
%
% if fixed_params=[NaN, NaN , NaN , NaN]        % or fixed_params=[]
% optimization of "min", "max", "x50" and "slope" (default)
%
% if fixed_params=[0, 1 , NaN , NaN]
% optimization of x50 and slope of a sigmoid of ranging from 0 to 1
%
%
% Additional information in the second output, STAT
% [param,stat]=sigm_fit(x,y,fixed_params,initial_params,plot_flag)
%
%
% Example:
% %% generate data vectors (x and y)
% fsigm = @(param,xval) param(1)+(param(2)-param(1))./(1+10.^((param(3)-xval)*param(4)))
% param=[0 1 5 1];  % "min", "max", "x50", "slope"
% x=0:0.1:10;
% y=fsigm(param,x) + 0.1*randn(size(x));
%
% %% standard parameter estimation
% [estimated_params]=sigm_fit(x,y)
%
% %% parameter estimation with forced 0.5 fixed min
% [estimated_params]=sigm_fit(x,y,[0.5 NaN NaN NaN])
%
% %% parameter estimation without plotting
% [estimated_params]=sigm_fit(x,y,[],[],0)
%
%
% Doubts, bugs: rpavao@gmail.com
% Downloaded from http://www.mathworks.com/matlabcentral/fileexchange/42641-sigmoid-logistic-curve-fit

% warning off

x=x(:);
y=y(:);

if nargin<=1 %fail
    fprintf('');
    help sigm_fit
    return
end

automatic_initial_params=[quantile(y,0.05) quantile(y,0.95) NaN 1];
if sum(y==quantile(y,0.5))==0
    temp=x(y==quantile(y(2:end),0.5));
else
    temp=x(y==quantile(y,0.5));
end
automatic_initial_params(3)=temp(1);

if nargin==2 %simplest valid input
    fixed_params=NaN(1,4);
    initial_params=automatic_initial_params;
    plot_flag=1;
end
if nargin==3
    initial_params=automatic_initial_params;
    plot_flag=1;
end
if nargin==4
    plot_flag=1;
end

if exist('fixed_params','var')
    if isempty(fixed_params)
        fixed_params=NaN(1,4);
    end
end
if exist('initial_params','var')
    if isempty(initial_params)
        initial_params=automatic_initial_params;
    end
end
if exist('plot_flag','var')
    if isempty(plot_flag)
        plot_flag=1;
    end
end

%p(1)=min; p(2)=max-min; p(3)=x50; p(4)=slope como em Y=Bottom + (Top-Bottom)/(1+10^((LogEC50-X)*HillSlope))
%f = @(p,x) p(1) + (p(2)-p(1)) ./ (1 + 10.^((p(3)-x)*p(4)));

f_str='f = @(param,xval)';
free_param_count=0;
bool_vec=NaN(1,4);
for i=1:4;
    if isnan(fixed_params(i))
        free_param_count=free_param_count+1;
        f_str=[f_str ' param(' num2str(free_param_count) ')'];
        bool_vec(i)=1;
    else
        f_str=[f_str ' ' num2str(fixed_params(i))];
        bool_vec(i)=0;
    end
    if i==1; f_str=[f_str ' + (']; end
    if i==2;
        if isnan(fixed_params(1))
            f_str=[f_str '-param(1) )./ (   1 + 10.^( ('];
        else
            f_str=[f_str '-' num2str(fixed_params(1)) ')./ (1 + 10.^(('];
        end
    end
    if i==3; f_str=[f_str ' - xval ) *']; end
    if i==4; f_str=[f_str ' )   );']; end
end

eval(f_str)

[BETA,RESID,J,COVB,MSE] = nlinfit(x,y,f,initial_params(bool_vec==1));
stat.param=BETA';

% confidence interval of the parameters
stat.paramCI = nlparci(BETA,RESID,'Jacobian',J);

% confidence interval of the estimation
[stat.ypred,delta] = nlpredci(f,x,BETA,RESID,'Covar',COVB);
stat.ypredlowerCI = stat.ypred - delta;
stat.ypredupperCI = stat.ypred + delta;

% plot(x,y,'ko') % observed data
% hold on
% plot(x,ypred,'k','LineWidth',2)
% plot(x,[lower,upper],'r--','LineWidth',1.5)

free_param_count=0;
for i=1:4;
    if isnan(fixed_params(i))
        free_param_count=free_param_count+1;
        param(i)=BETA(free_param_count);
    else
        param(i)=fixed_params(i);
    end
end

if plot_flag==1
    x_vector=min(x):(max(x)-min(x))/100:max(x);
    figure;
    plot(x,y,'k.',x_vector,f(param(isnan(fixed_params)),x_vector),'r-')
    xlim([min(x) max(x)])
end
end
%% FUNCTION
function cmap = ColorMapDerek(start_col,end_col,invert)
    cmap = [linspace(start_col(1),end_col(1),256)', linspace(start_col(2),end_col(2),256)', linspace(start_col(3),end_col(3),256)'];
    if strcmp(invert,'yes'); cmap = 1-cmap; end
end